import os
import weakref
from itertools import count

from fsspec import AbstractFileSystem
from zkl_pyutils_fsspec import FsLike, resolve_child_fs, resolve_fs


from .str_embeddings_writer import StrEmbeddingsWriter



class OmicsDatasetWriter:
    def __init__(self, fs: FsLike):
        fs = resolve_fs(fs)
        fs.mkdir("", exists_ok=True)
        write_dataset_script(fs)
        self._fs = fs
        self._opened_writers = weakref.WeakSet()
        self._embedder = None




    def str_embeddings_writer(self, name: str | None = None):
        name = make_columns_name(self._fs, name, "str_embeddings")
        fs = resolve_child_fs(self._fs, name)
        writer = StrEmbeddingsWriter(fs, chunk_size=4 * 2 ** 30)
        self._opened_writers.add(writer)
        return writer


    def flush(self):
        for writer in self._opened_writers:
            writer.flush()

    def close(self):
        for writer in self._opened_writers:
            writer.close()
        if self._embedder is not None:
            self._embedder.close()
            self._embedder = None

    def __del__(self):
        self.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


dataset_script_file_path = os.path.join(os.path.dirname(__file__), "omics_dataset.py")


def write_dataset_script(fs: AbstractFileSystem):
    with open(dataset_script_file_path, "rt") as fp:
        dataset_script = fp.read()
    with fs.open("dataset.py", "wt") as fp:
        fp.write(dataset_script)


def make_columns_name(fs: AbstractFileSystem, name: str | None, default: str):
    if name is not None:
        if fs.exists(name):
            raise FileExistsError
        return name
    if not fs.exists(default):
        return default
    for i in count(start=1):
        candidate = f"{default}_{i}"
        if not fs.exists(candidate):
            return candidate
    raise RuntimeError("Failed to make the name!")
